{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/NeuromatchAcademy/course-content/blob/master/tutorials/W3D2_HiddenDynamics/student/W3D2_Tutorial2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Neuromatch Academy: Week 3, Day 2, Tutorial 2\n",
    "\n",
    "# Hidden Markov Model\n",
    "\n",
    "__Content creators:__ Yicheng Fei with help from Jesse Livezey\n",
    "\n",
    "__Content reviewers:__ John Butler, Matt Krause, Meenakshi Khosla, Spiros Chavlis, Michael Waskom\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial objectives\n",
    "\n",
    "The world around us is often changing state over time, but we may only have access to these states through noisy sensory measurements. Similarly, organisms and neural systems often are thought to transition between a set of discrete states (up/down states, sleep/wake, etc.) which may only be indirectly observable through their impact on neural activity. Hidden Markov Models are a class of models that allow us to reason about the dynamics of a set of unobserved states that lead to the changing sensory inputs or data we observe.\n",
    "\n",
    "In this notebook, we'll first simulate a Hidden Markov Model and observe how changing the transition probability and observation noise impact what the samples look like. Then we'll look at how uncertainty increases as we make future predictions without evidence (from observations) and how to gain information from the observations.\n",
    "The HMM model we use in the first part of the tutorial will have a binary latent variable $s_t \\in \\{0,1\\}$ that switches randomly between the two states, and a 1D Gaussian emission model $m_t|s_t \\sim \\mathcal{N}(\\mu_{s_t},\\sigma^2_{s_t})$ that provides evidence about the current state. You will learn how to:\n",
    "\n",
    "* Build an HMM in Python and generate sample data.\n",
    "* Calculate how predictive probabilities propagates in a Markov Chain with no evidence.\n",
    "* Combine new evidence and prediction from past evidence  to estimate latent states.\n",
    "\n",
    "---\n",
    "\n",
    "There is an an **optional** part for you to get a sense of how to perform parameter estimation of an HMM using the EM algorithm. **We encourage you to do these bonus exercises only _after_ you complete the core material in Tutorials 3 and 4.**\n",
    "\n",
    "In the optional part, you will implement an HMM of a network of Poisson spiking neurons mentioned in today's intro and:\n",
    "\n",
    "* Implement the forward-backward algorithm\n",
    "* Complete the E-step and M-step \n",
    "* Learn parameters for the example problem using the EM algorithm\n",
    "* Get an intuition of how the EM algorithm monotonically increases data likelihood"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:26.786248Z",
     "iopub.status.busy": "2021-06-03T14:00:26.785502Z",
     "iopub.status.idle": "2021-06-03T14:00:26.821883Z",
     "shell.execute_reply": "2021-06-03T14:00:26.822432Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Video 1: Introduction\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"ceQXN0OUaFo\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:26.829442Z",
     "iopub.status.busy": "2021-06-03T14:00:26.828209Z",
     "iopub.status.idle": "2021-06-03T14:00:29.197671Z",
     "shell.execute_reply": "2021-06-03T14:00:29.198140Z"
    }
   },
   "outputs": [],
   "source": [
    "!pip install hmmlearn --quiet\n",
    "\n",
    "import numpy as np\n",
    "from scipy import stats\n",
    "from scipy.optimize import linear_sum_assignment\n",
    "from hmmlearn import hmm\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import patches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:29.204462Z",
     "iopub.status.busy": "2021-06-03T14:00:29.202585Z",
     "iopub.status.idle": "2021-06-03T14:00:29.391865Z",
     "shell.execute_reply": "2021-06-03T14:00:29.391244Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Figure Settings\n",
    "import ipywidgets as widgets       # interactive display\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "plt.style.use(\"https://raw.githubusercontent.com/NeuromatchAcademy/course-content/NMA2020/nma.mplstyle\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:29.413223Z",
     "iopub.status.busy": "2021-06-03T14:00:29.405020Z",
     "iopub.status.idle": "2021-06-03T14:00:29.416930Z",
     "shell.execute_reply": "2021-06-03T14:00:29.416465Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Helper functions\n",
    "\n",
    "def plot_hmm1(model, states, observations):\n",
    "  \"\"\"Plots HMM states and observations for 1d states and observations.\n",
    "\n",
    "  Args:\n",
    "    model (hmmlearn model):               hmmlearn model used to get state means.\n",
    "    states (numpy array of floats):       Samples of the states.\n",
    "    observations (numpy array of floats): Samples of the states.\n",
    "  \"\"\"\n",
    "\n",
    "  nsteps = states.size\n",
    "  fig, ax1 = plt.subplots()\n",
    "  states_forplot = list(map(lambda s: model.means_[s], states))\n",
    "  ax1.step(np.arange(nstep), states_forplot, \"--\", where=\"mid\", alpha=1.0, c=\"green\")\n",
    "  ax1.set_xlabel(\"Time\")\n",
    "  ax1.set_ylabel(\"Latent State\", c=\"green\")\n",
    "  ax1.set_yticks([-1, 1])\n",
    "  ax1.set_yticklabels([\"State 1\", \"State 0\"])\n",
    "\n",
    "  ax2 = ax1.twinx()\n",
    "  ax2.plot(np.arange(nstep), observations.flatten(), c=\"blue\")\n",
    "  ax2.set_ylabel(\"Observations\", c=\"blue\")\n",
    "  ax1.set_ylim(ax2.get_ylim())\n",
    "  plt.show(fig)\n",
    "\n",
    "\n",
    "def plot_marginal_seq(predictive_probs, switch_prob):\n",
    "  \"\"\"Plots the sequence of marginal predictive distributions.\n",
    "\n",
    "    Args:\n",
    "      predictive_probs (list of numpy vectors): sequence of predictive probability vectors\n",
    "      switch_prob (float):                      Probability of switching states.\n",
    "  \"\"\"\n",
    "  T = len(predictive_probs)\n",
    "  prob_0 = [p_vec[0] for p_vec in predictive_probs]\n",
    "  prob_1 = [p_vec[1] for p_vec in predictive_probs]\n",
    "  fig, ax = plt.subplots()\n",
    "  ax.plot(np.arange(T), prob_0, color=\"orange\")\n",
    "  ax.plot(np.arange(T), prob_1, color=\"blue\")\n",
    "  ax.legend([\n",
    "    \"prob in state 0\", \"prob in state 1\"\n",
    "  ])\n",
    "  ax.text(T/2, 0.05, \"switching probability={}\".format(switch_prob), fontsize=12,\n",
    "          bbox=dict(boxstyle=\"round\", facecolor=\"wheat\", alpha=0.6))\n",
    "  ax.set_xlabel(\"Time\")\n",
    "  ax.set_ylabel(\"Probability\")\n",
    "  ax.set_title(\"Forgetting curve in a changing world\")\n",
    "  plt.show(fig)\n",
    "\n",
    "def plot_evidence_vs_noevidence(posterior_matrix, predictive_probs):\n",
    "  \"\"\"Plots the average posterior probabilities with evidence v.s. no evidence\n",
    "\n",
    "  Args:\n",
    "    posterior_matrix: (2d numpy array of floats): The posterior probabilities in state 1 from evidence (samples, time)\n",
    "    predictive_probs (numpy array of floats):  Predictive probabilities in state 1 without evidence\n",
    "  \"\"\"\n",
    "  nsample, T = posterior_matrix.shape\n",
    "  posterior_mean = posterior_matrix.mean(axis=0)\n",
    "  fig, ax = plt.subplots(1)\n",
    "  ax.plot([0.0, T],[0.5, 0.5], color=\"red\", linestyle=\"dashed\")\n",
    "  ax.plot(np.arange(T), predictive_probs, c=\"orange\", linewidth=2, label=\"No evidence\")\n",
    "  ax.scatter(np.tile(np.arange(T), (nsample, 1)), posterior_matrix, s=0.8, c=\"green\", alpha=0.3, label=\"With evidence(Sample)\")\n",
    "  ax.plot(np.arange(T), posterior_mean, c='green', linewidth=2, label=\"With evidence(Average)\")\n",
    "  ax.legend()\n",
    "  ax.set_yticks([0.0, 0.25, 0.5, 0.75, 1.0])\n",
    "  ax.set_xlabel(\"Time\")\n",
    "  ax.set_ylabel(\"Probability in State 0\")\n",
    "  ax.set_title(\"Gain confidence with evidence\")\n",
    "  plt.show(fig)\n",
    "\n",
    "def simulate_forward_inference(model, T, data=None):\n",
    "  \"\"\"\n",
    "  Given HMM `model`, calculate posterior marginal predictions of x_t for T-1 time steps ahead based on\n",
    "  evidence `data`. If `data` is not give, generate a sequence of observations from first component.\n",
    "\n",
    "  Args:\n",
    "    model (GaussianHMM instance): the HMM\n",
    "    T (int): length of returned array\n",
    "\n",
    "  Returns:\n",
    "    predictive_state1: predictive probabilities in first state w.r.t no evidence\n",
    "    posterior_state1: posterior probabilities in first state w.r.t evidence\n",
    "  \"\"\"\n",
    "\n",
    "  # First re-calculate hte predictive probabilities without evidence\n",
    "  predictive_probs = simulate_prediction_only(model, T)\n",
    "  # Generate an observation trajectory condtioned on that latent state x is always 1\n",
    "  if data is not None:\n",
    "    Y = data\n",
    "  else:\n",
    "    Y = np.asarray([model._generate_sample_from_state(0) for _ in range(T)])\n",
    "  # Calculate marginal for each latent state x_t\n",
    "  pt = np.exp(model._compute_log_likelihood(Y[[0]])) * model.startprob_\n",
    "  pt /= np.sum(pt)\n",
    "  posterior_probs = np.zeros((T, pt.size))\n",
    "  posterior_probs[0] = pt\n",
    "\n",
    "  for t in range(1, T):\n",
    "    posterior = one_step_update(model, posterior_probs[t-1], Y[[t]])\n",
    "    # normalize and add to the list\n",
    "    posterior /= np.sum(posterior)\n",
    "    posterior_probs[t] = posterior\n",
    "  posterior_state1 = np.asarray([p[0] for p in posterior_probs])\n",
    "  predictive_state1 = np.asarray([p[0] for p in predictive_probs])\n",
    "  return predictive_state1, posterior_state1\n",
    "\n",
    "def plot_forward_inference(model, states, observations, states_inferred):\n",
    "    \"\"\"Plot ground truth state sequence with noisy observations, and ground truth states v.s. inferred ones\n",
    "\n",
    "        Args:\n",
    "            model (instance of hmmlearn.GaussianHMM): an instance of HMM\n",
    "            states (numpy vector): vector of 0 or 1(int or Bool), the sequences of true latent states\n",
    "            observations (numpy vector of numpy vector): the un-flattened Gaussian observations at each time point, element has size (1,)\n",
    "            states_inferred (numpy vector): vector of 0 or 1(int or Bool), the sequences of inferred latent states\n",
    "    \"\"\"\n",
    "    plot_hmm1(model, states, observations)\n",
    "    fig, ax = plt.subplots()\n",
    "    # state 0 has larger mean\n",
    "    ax.step(np.arange(nstep), 1-states, color=\"green\", label=\"Ground Truth\")\n",
    "    ax.step(np.arange(nstep), 1-states_inferred, linestyle=\"dashed\", color=\"orange\", label=\"Inferred\")\n",
    "    ax.legend()\n",
    "    ax.set_title(\"Infer latent states from data\")\n",
    "    ax.set_xlabel(\"Time\")\n",
    "    ax.set_ylabel(\"Latent State\")\n",
    "    ax.set_yticks([0, 1])\n",
    "    ax.set_yticklabels([\"State 1\", \"State 0\"])\n",
    "    plt.show(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Section 1: Binary HMM with Gaussian observations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:29.422675Z",
     "iopub.status.busy": "2021-06-03T14:00:29.421929Z",
     "iopub.status.idle": "2021-06-03T14:00:29.454065Z",
     "shell.execute_reply": "2021-06-03T14:00:29.453040Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Video 2: Simulating a binary HMM with Gaussian observations\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"7cTnoe6Xt80\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In contrast with the Sequential Probability Ratio Test, the latent state in an HMM is no longer fixed over time. Instead, it can probabilistically switch or jump to a different state at each time step. However, the  time dependence of states at different times is simple: the probability of the state at time $t$ is wholely determined by the state at time $t-1$. This is called called the **Markov property** and the dependency of the whole state sequence $\\{s_1,...,s_t\\}$ can be described by a chain structure called a Markov Chain:\n",
    "\n",
    "<img src=\" https://github.com/NeuromatchAcademy/course-content/blob/master/tutorials/W2D3_DecisionMaking/static/W2D3_Tutorial2_markov_chain_diagram.png?raw=true\" alt=\"Markov chain drawing\" width=\"400\"/>\n",
    "\n",
    "(Please note that this should be s in diagram above, it will be changed)\n",
    "\n",
    "**Markov model for latent dynamics**\n",
    "\n",
    "Here we will reuse the switching process or telegraph process you saw in a previous tutorial. Quantitatively, the probability of switching to state $s_t=j$ from the previous state $s_{t-1}=i$ is a conditional probability distribution $p(s_t=j|s_{t-1}=i)$.\n",
    "\n",
    "Since the states are binary, we can represent the probability of the current state as a 2-dimensional vector $p(s=i)=p_{i}$ (or, including time, as $p(s_t=i)=p_{ti}$), and can represent the transition probability as a 2$\\times$2 matrix $A_{ij}$. This is a convenient representation for coding. We can then use this representation to update the probabilities over time following the Markov process.\n",
    "$$p(s_t=j) = \\sum_{i} p(s_t=j|s_{t-1}=i)p(s_{t-1}=i)$$\n",
    "or equivalently\n",
    "$$p_{tj}=\\sum_j A_{ij} p_{(t-1)i} \\tag{1}$$\n",
    "or, using vectors, $p_t=Ap_{t-1}$. Note that here $A_{ij}$ represents the transition probability to switch **FROM state $i$ TO state $j$** at next time step.\n",
    "\n",
    "**Measurements**\n",
    "\n",
    "In a _Hidden_ Markov model, we cannot directly observe the latent states $s_t$. What we can observe instead is a noisy measurement $m_t$ generated from $s_t$.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Coding Exercise 1: Simulate a binary HMM with Gaussian observations\n",
    "\n",
    "In this exercise, you will use the package `hmmlearn` to implement a two-state HMM with Gaussian measurements (sometimes called emissions). Your HMM will start in State 0 and transition between states (both $0 \\rightarrow 1$ and $1 \\rightarrow 0$) with probability `switch_prob`. Each state emits observations drawn from a Gaussian with $\\mu = 1$ for State 0 and $\\mu = -1$ for State 1. The variance of both states is fixed at `noise_level`.\n",
    "\n",
    "Please familiarize yourself with the code and complete the following exercises in the next cell. You will need to:\n",
    "\n",
    "1. To implement the state transitions, complete the transition matrix  `transmat_` (i.e., $A_{ij}$) in the code below. \n",
    "\\begin{equation*}\n",
    "A_{i,j} = \n",
    "\\begin{pmatrix}\n",
    "p_{\\rm stay} & p_{\\rm switch} \\\\\n",
    "p_{\\rm switch} & p_{\\rm stay} \\\\\n",
    "\\end{pmatrix}\n",
    "\\end{equation*}\n",
    "with $p_{\\rm stay} = 1 - p_{\\rm switch}$. \n",
    "\n",
    "2. The *hidden* part of HMM means that we do not directly output the current state $s_t$, but instead observe a noisy emission $m_t | s_t$, here generated by a Gaussian. The  means have already been filled in for you, but you must complete the covariance matrix `covars_`. Set each state's observation variance to `noise_level`. In the code, the required shape given below is $2\\times 1\\times 1$, for two $1\\times 1$ covariances which are really scalar variances. This seems like a weird shape for storing two numbers, but it makes things easier for the rest of the code. \n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:29.462575Z",
     "iopub.status.busy": "2021-06-03T14:00:29.460016Z",
     "iopub.status.idle": "2021-06-03T14:00:29.584905Z",
     "shell.execute_reply": "2021-06-03T14:00:29.583709Z"
    }
   },
   "outputs": [],
   "source": [
    "def create_HMM(switch_prob=0.1, noise_level=1e-8, startprob=[1.0, 0.0]):\n",
    "    \"\"\"Create an HMM with binary state variable and 1D Gaussian observations\n",
    "    The probability to switch to the other state is `switch_prob`. Two\n",
    "    observation models have mean 1.0 and -1.0 respectively. `noise_level`\n",
    "    specifies the standard deviation of the observation models.\n",
    "\n",
    "    Args:\n",
    "        switch_prob (float): probability to jump to the other state\n",
    "        noise_level (float): standard deviation of observation models. Same for\n",
    "        two components\n",
    "\n",
    "    Returns:\n",
    "        model (hmm.GaussianHMM instance): the described HMM\n",
    "    \"\"\"\n",
    "    ############################################################################\n",
    "    # Insert your code here to:\n",
    "    #    * Create the transition matrix, `transmat_` so that the odds of\n",
    "    #      switching is `switch_prob`\n",
    "    #\t\t* Set the observation model variances, `covars_`, to `noise_level`\n",
    "    raise NotImplementedError(\"`create_HMM` is incomplete\")\n",
    "    ############################################################################\n",
    "    n_components = 2\n",
    "\n",
    "    # Initialize model\n",
    "    model = hmm.GaussianHMM(n_components=n_components, covariance_type=\"full\")\n",
    "    model.startprob_ = np.asarray(startprob)\n",
    "\n",
    "    # Make transition matrix, should be shape (2, 2), i.e., a transition matrix for 2 states\n",
    "    model.transmat_ = ...\n",
    "\n",
    "    # Create means\n",
    "    model.means_ = np.array([[1.0], [-1.0]])\n",
    "\n",
    "    # Create covariance matrices, should be shape (2, 1, 1), i.e., 2 1x1 covariance matrices\n",
    "    model.covars_ = ...\n",
    "\n",
    "    model.sample(1)\n",
    "\n",
    "    return model\n",
    "\n",
    "\n",
    "# Set random seed\n",
    "np.random.seed(101)\n",
    "\n",
    "# Number of steps\n",
    "nstep = 50\n",
    "\n",
    "# Create HMM\n",
    "model = create_HMM()\n",
    "\n",
    "# Sample from HMM\n",
    "observations, states = model.sample(nstep)\n",
    "\n",
    "# Visualize\n",
    "plot_hmm1(model, states, observations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:29.597780Z",
     "iopub.status.busy": "2021-06-03T14:00:29.596167Z",
     "iopub.status.idle": "2021-06-03T14:00:29.973400Z",
     "shell.execute_reply": "2021-06-03T14:00:29.972479Z"
    }
   },
   "source": [
    "[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D2_HiddenDynamics/solutions/W3D2_Tutorial2_Solution_9e86d4ae.py)\n",
    "\n",
    "*Example output:*\n",
    "\n",
    "<img alt='Solution hint' align='left' width=560 height=416 src=https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/tutorials/W3D2_HiddenDynamics/static/W3D2_Tutorial2_Solution_9e86d4ae_0.png>\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interactive Demo 1: Binary HMM\n",
    "\n",
    "In the demo below, we simulate a similar HMM and plot. You can change the probability of switching states and the noise level.\n",
    "\n",
    "First, think and discuss these questions.\n",
    "\n",
    "\n",
    "1.   What will happen if the switching probability is zero? What about if it's one?\n",
    "2.   What will happen with high noise? Low?\n",
    "\n",
    "\n",
    "\n",
    "Then, play with the demo to see if you were correct or not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:30.000781Z",
     "iopub.status.busy": "2021-06-03T14:00:29.985257Z",
     "iopub.status.idle": "2021-06-03T14:00:30.356135Z",
     "shell.execute_reply": "2021-06-03T14:00:30.355570Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "\n",
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "np.random.seed(101)\n",
    "nstep = 100\n",
    "\n",
    "@widgets.interact\n",
    "def plot(switch_prob=(0., 1, .01), log10_noise_level=(-8., 1., .01)):\n",
    "  model = create_HMM(switch_prob=switch_prob,\n",
    "                     noise_level=10.**log10_noise_level)\n",
    "\n",
    "  observations, states = model.sample(nstep)\n",
    "  observations = observations.flatten()\n",
    "  plot_hmm1(model, states, observations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:30.363834Z",
     "iopub.status.busy": "2021-06-03T14:00:30.363225Z",
     "iopub.status.idle": "2021-06-03T14:00:30.366128Z",
     "shell.execute_reply": "2021-06-03T14:00:30.366547Z"
    }
   },
   "source": [
    "[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D2_HiddenDynamics/solutions/W3D2_Tutorial2_Solution_41ccf32d.py)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "**Harkening** back to our fishing example, you can imagine that the time series you measure is related to the number of fish caught at different times as the school of fish move from left to right. Or you could envision it as the voltage across a membrane affected by an ion channel in two states, open and closed. Or it could represent EEG frequency measurements as the brain moves between sleep states. What phenomena can you imagine modeling with these HMMs?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# Section 2: Forgetting information and gaining confidence with a known initial state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 2.1: Forgetting information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:30.372887Z",
     "iopub.status.busy": "2021-06-03T14:00:30.372204Z",
     "iopub.status.idle": "2021-06-03T14:00:30.403094Z",
     "shell.execute_reply": "2021-06-03T14:00:30.403679Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Video 3: Forgetting in a changing world\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"pRRo_L-n8nc\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Coding Exercise 2: Forgetting in a changing world\n",
    "\n",
    "Since the world (latent state) is changing over time, even if we know for sure that we are in state 0 at some time, we will be more and more uncertain that we'll remain in state 0 as time goes. In other words, when we try to make predictions of future states in a Markov Chain based on our current knowledge without future evidence, the influence of current state will decay over time.\n",
    "\n",
    "In this exercise, we'll inspect how we \"forget\" the current state information when predicting future states without any observation.\n",
    "\n",
    "Using the model you just defined, let's now make some predictions about $s_t$ given that we know $s_0=0$ for sure. We've already imposed this assumption by setting prior probabilities of $p(s_0)$ to $[1,0]$ earlier. \n",
    "\n",
    "1. Complete the code in function `markov_forward` to calculate the predictive marginal distribution at next time step using `p_next = A.T @ p_current` \n",
    "\n",
    "3. Take a look at function `simulate_prediction_only` and understand how the predictive distribution propagates along the Markov chain\n",
    "\n",
    "4. Using our provided code, plot the predictive probabilities as a function of time\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:30.409707Z",
     "iopub.status.busy": "2021-06-03T14:00:30.408279Z",
     "iopub.status.idle": "2021-06-03T14:00:30.410357Z",
     "shell.execute_reply": "2021-06-03T14:00:30.411134Z"
    }
   },
   "outputs": [],
   "source": [
    "# @markdown Execute this cell to enable the function simulate_prediction_only\n",
    "def simulate_prediction_only(model, nstep):\n",
    "  \"\"\"\n",
    "  Simulate the diffusion of HMM with no observations\n",
    "\n",
    "  Args:\n",
    "    model (hmm.GaussianHMM instance): the HMM instance\n",
    "    nstep (int): total number of time steps to simulate(include initial time)\n",
    "\n",
    "  Returns:\n",
    "    predictive_probs (list of numpy vector): the list of marginal probabilities\n",
    "  \"\"\"\n",
    "  entropy_list = []\n",
    "  predictive_probs = []\n",
    "  prob = model.startprob_\n",
    "  for i in range(nstep):\n",
    "    # calculate entropy\n",
    "    predictive_probs.append(prob)\n",
    "    # one step forward\n",
    "    prob = markov_forward(prob, model.transmat_)\n",
    "  return predictive_probs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:30.416983Z",
     "iopub.status.busy": "2021-06-03T14:00:30.416443Z",
     "iopub.status.idle": "2021-06-03T14:00:30.431954Z",
     "shell.execute_reply": "2021-06-03T14:00:30.431386Z"
    }
   },
   "outputs": [],
   "source": [
    "def markov_forward(p0, A):\n",
    "  \"\"\"Calculate the forward predictive distribution in a discrete Markov chain\n",
    "\n",
    "  Args:\n",
    "    p0 (numpy vector): a discrete probability vector\n",
    "    A (numpy matrix): the transition matrix, A[i,j] means the prob. to\n",
    "    switch FROM i TO j\n",
    "\n",
    "  Returns:\n",
    "    p1 (numpy vector): the predictive probabilities in next time step\n",
    "  \"\"\"\n",
    "  ############################################################################\n",
    "  # Insert your code here to:\n",
    "  #      Compute the marginal distribution of Markov chain in next time step\n",
    "  #      Hint: use matrix multiply and be careful about the index orders\n",
    "  raise NotImplementedError(\"function `markov_forward` incomplete\")\n",
    "  ############################################################################\n",
    "  p1 = ...\n",
    "  return p1\n",
    "\n",
    "\n",
    "# Set random seed\n",
    "np.random.seed(101)\n",
    "\n",
    "# Set parameters of HMM\n",
    "T = 100\n",
    "switch_prob = 0.1\n",
    "noise_level = 2.0\n",
    "\n",
    "# Create HMM\n",
    "model = create_HMM(switch_prob=switch_prob, noise_level=noise_level)\n",
    "\n",
    "# Get predictive probabilities\n",
    "predictive_probs = simulate_prediction_only(model, T)\n",
    "\n",
    "# Visualize\n",
    "plot_marginal_seq(predictive_probs, switch_prob)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:30.481462Z",
     "iopub.status.busy": "2021-06-03T14:00:30.460426Z",
     "iopub.status.idle": "2021-06-03T14:00:30.770704Z",
     "shell.execute_reply": "2021-06-03T14:00:30.771204Z"
    }
   },
   "source": [
    "[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D2_HiddenDynamics/solutions/W3D2_Tutorial2_Solution_61403d8f.py)\n",
    "\n",
    "*Example output:*\n",
    "\n",
    "<img alt='Solution hint' align='left' width=560 height=416 src=https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/tutorials/W3D2_HiddenDynamics/static/W3D2_Tutorial2_Solution_61403d8f_0.png>\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interactive Demo 2: Forgetting\n",
    "\n",
    "In the following demo, we look at the same visualization but you can play with the probability of switching states, using the slider. \n",
    "\n",
    "1.  Do you forget more quickly with low or high switching probability?\n",
    "\n",
    "2.  How does the curve look at when `prob_switch` $>0.5?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:30.794459Z",
     "iopub.status.busy": "2021-06-03T14:00:30.793784Z",
     "iopub.status.idle": "2021-06-03T14:00:31.081200Z",
     "shell.execute_reply": "2021-06-03T14:00:31.080553Z"
    }
   },
   "outputs": [],
   "source": [
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "np.random.seed(101)\n",
    "T = 100\n",
    "noise_level = 0.5\n",
    "\n",
    "@widgets.interact\n",
    "def plot(switch_prob=(0.01, .99, .01)):\n",
    "  model = create_HMM(switch_prob=switch_prob, noise_level=noise_level)\n",
    "  predictive_probs = simulate_prediction_only(model, T)\n",
    "  plot_marginal_seq(predictive_probs, switch_prob)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:31.086430Z",
     "iopub.status.busy": "2021-06-03T14:00:31.085835Z",
     "iopub.status.idle": "2021-06-03T14:00:31.091338Z",
     "shell.execute_reply": "2021-06-03T14:00:31.091927Z"
    }
   },
   "source": [
    "[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D2_HiddenDynamics/solutions/W3D2_Tutorial2_Solution_d90e17f2.py)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 2.2: Gaining confidence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:31.097195Z",
     "iopub.status.busy": "2021-06-03T14:00:31.096468Z",
     "iopub.status.idle": "2021-06-03T14:00:31.129967Z",
     "shell.execute_reply": "2021-06-03T14:00:31.130413Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Video 4: Gain confidence from evidence\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"dDjoxUxMgC0\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As shown in Exercise 2, you lose information and increase uncertainty exponentially when predicting into future in the absence of further evidence, because the state randomly diffuses from its last known value. However the HMM also generates a measurement $m_t$ at each time step, and we can use this evidence to improve our state estimate.\n",
    "\n",
    "<img src=\"https://github.com/NeuromatchAcademy/course-content/blob/master/tutorials/W2D3_DecisionMaking/static/W2D3_Tutorial2_HMM_diagram.png?raw=true\" alt=\"HMM drawing\" width=\"400\"/>\n",
    "\n",
    "(We will UPDATE FIGURE with $x,y\\to s,m$**)\n",
    "\n",
    "Now let's incorporate evidence into our inference. In this exercise we will calculate the **posterior marginal** distribution $p(s_t|m_{1:t})$, ie the marginal probability of the current state given the entire history of measurements. This is a crucial computation, and it is tractable because of the simple structure of the HMM.\n",
    "\n",
    "We compute this probability recursively. Suppose we know the posterior marginal probability for the previous time step, $p(s_{t-1}|m_{1:t-1})$. Now we receive a new measurement $m_t$. From Bayes' rule and the Markov property, we can calculate $p(s_{t}|m_{1:t})$. We do this in two steps.\n",
    "\n",
    "First, we make a prediction for $s_t$ given our previous knowledge. We can say \"yesterday's posterior becomes today's prior,\" where the Markov transition matrix accounts for the change from $t$ to $t-1$, as in the last exercise. This gives us the prior probability\n",
    "$$p(s_t|m_{1:t-1})=\\sum_{s_{t-1}}p(s_t|s_{t-1})p(s_{t-1}|m_{1:t-1}) \\tag{2}$$\n",
    "Observe that the history $m_{1:t-1}$ does not yet include the new measurement at time $t$.\n",
    "\n",
    "Second, we use the usual Bayesian inference to incorporate this new evidence, multiplying our prior $p(s_t|m_{1:t-1})$ times our likelihood $p(m_t|s_t)$, to obtain\n",
    "$$p(s_t|m_{1:t})\\propto p(s_t|m_{1:t-1})p(s_{t}|m_t) \\tag{3}$$.\n",
    "\n",
    "Putting (2) and (3) together, we obtain the forward recursion equation for a Hidden Markov Model,\n",
    "$$p(s_t|m_{1:t-1})\\propto p(m_t|s_t)\\sum_{s_{t-1}}p(s_t|s_{t-1})p(s_{t-1}|m_{1:t-1}) \\tag{4}$$\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Coding Exercise 3: Gain confidence from evidence\n",
    "\n",
    "1. Create a model with switching probability $0.1$ and noise level $0.5$ using function `create_HMM(switch_prob, noise_level)`\n",
    "\n",
    "2. Complete the code to calculate marginal posterior distribution $p(s_t|m_{1:t-1})$ at time $t$ from last posterior $p(s_{t-1}|m_{1:t-1})$ at time $t-1$ \n",
    "    - Calculate the predictive distribution $p(s_t =j|m_{1:t-1})=\\sum_i A_{ij} p(s_{t-1}=i|m_{1:t-1})$\n",
    "    - Calculate the likelihood of new data under each component using `exp(model._compute_log_likelihood(yt))`\n",
    "    - Multiply likelihood and prediction element-wise and normalize over two components to get the new posterior probabilities\n",
    "\n",
    "3. Using provided code, plot the average posterior probabilities over time due to evidence together with predictive probabilities without evidence\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:31.138893Z",
     "iopub.status.busy": "2021-06-03T14:00:31.137522Z",
     "iopub.status.idle": "2021-06-03T14:00:31.173934Z",
     "shell.execute_reply": "2021-06-03T14:00:31.172255Z"
    }
   },
   "outputs": [],
   "source": [
    "def one_step_update(model, posterior_tm1, Y_t):\n",
    "  \"\"\"Given a HMM model, calculate the one-time-step updates to the posterior.\n",
    "\n",
    "  Args:\n",
    "    model (GaussianHMM instance): the HMM\n",
    "    posterior_tm1 (numpy array): Posterior at `t-1`\n",
    "    Y_t (numpy array): Observation at `t`\n",
    "\n",
    "    Returns:\n",
    "    posterior_t (numpy array): Posterior at `t`\n",
    "  \"\"\"\n",
    "  ##############################################################################\n",
    "  # Insert your code here to:\n",
    "  #      1. Calculate the predicted state given the previous\n",
    "  #      estimate (`posterior_tm1`). Note that `model.transmat_` is equvalent\n",
    "  #      to `A.T`, not `A`.\n",
    "  #      2. Using `model._compute_log_likelihood()`, calculate the likelihood\n",
    "  #      given `Y_t`.\n",
    "  raise NotImplementedError(\"`one_step_update` is incomplete\")\n",
    "  ##############################################################################\n",
    "  prediction = ... @ posterior_tm1\n",
    "  likelihood = np.exp(...)\n",
    "  posterior_t = prediction * likelihood\n",
    "  return posterior_t\n",
    "\n",
    "\n",
    "np.random.seed(101)\n",
    "switch_prob = 0.1\n",
    "noise_level = 0.5\n",
    "nsample = 50\n",
    "T = 160\n",
    "\n",
    "model = create_HMM(switch_prob, noise_level)\n",
    "\n",
    "posterior_list = []\n",
    "for i in range(nsample):\n",
    "  predictive_probs, posterior_probs = simulate_forward_inference(model, T)\n",
    "  posterior_list.append(posterior_probs)\n",
    "posterior_matrix = np.asarray(posterior_list)\n",
    "\n",
    "plot_evidence_vs_noevidence(posterior_matrix, predictive_probs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:31.187392Z",
     "iopub.status.busy": "2021-06-03T14:00:31.184308Z",
     "iopub.status.idle": "2021-06-03T14:00:34.581436Z",
     "shell.execute_reply": "2021-06-03T14:00:34.580971Z"
    }
   },
   "source": [
    "[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D2_HiddenDynamics/solutions/W3D2_Tutorial2_Solution_0601fea5.py)\n",
    "\n",
    "*Example output:*\n",
    "\n",
    "<img alt='Solution hint' align='left' width=560 height=416 src=https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/tutorials/W3D2_HiddenDynamics/static/W3D2_Tutorial2_Solution_0601fea5_0.png>\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you've got a plot of information loss due to diffusion together with the information recovered/uncertainty reduced due to evidence. The difference between the former and the latter is the amount of uncertainty that still remains because of observation noise, as we'll see in the next demo. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interactive Demo 2.2: Noise in confidence\n",
    "In this demo, you can adjust the switch probability and noise level and observe how information gain changes with signal-to-noise ratio and/or switch probability.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:34.626447Z",
     "iopub.status.busy": "2021-06-03T14:00:34.589333Z",
     "iopub.status.idle": "2021-06-03T14:00:40.798167Z",
     "shell.execute_reply": "2021-06-03T14:00:40.797629Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "\n",
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "np.random.seed(101)\n",
    "@widgets.interact\n",
    "def plot(switch_prob=(0.01, .99, .01), noise_level=(.1, 3, .05),\n",
    "         nsample=(5, 200, 5), T=(20, 300, 5)):\n",
    "  model = create_HMM(switch_prob, noise_level)\n",
    "  posterior_list = []\n",
    "  for i in range(nsample):\n",
    "    predictive_probs, posterior_probs = simulate_forward_inference(model, T)\n",
    "    posterior_list.append(posterior_probs)\n",
    "  posterior_matrix = np.asarray(posterior_list)\n",
    "  plot_evidence_vs_noevidence(posterior_matrix, predictive_probs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# Section 3: Inference in a dynamic world"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Coding Exercise 4: Forward inference of HMM\n",
    "\n",
    "If you set `switch_prob` or `noise_level` to be large in the last exercise, you will observe that some sample inference dots fall below 0.5. This means we are making false inferences about which latent state we are in. \n",
    "\n",
    "In this exercise, let's make a forward inference of a random state sequence rather than a constant one by observing its noisy Gaussian outputs. Different from Exercise 1, here we assume we know the switching probability but don't know the prior (`startprob_`).\n",
    "\n",
    "1. Build a HMM with prior probabilities= $(0.5,0.5)$, switching probability=$0.1$, and noise level=$1.0$ by calling `create_HMM(switch_prob, noise_level, startprob)`\n",
    "2. Generate a sample sequence along with observations by calling `model.sample(nstep)`, and use our provided code to visualize the latent trajectory and observations\n",
    "3. Calculate posterior probabilities given data by calling `simulate_forward_inference(model, nstep, observations)`, and make inference of latent states by picking the component with larger posterior probability\n",
    "4. Use our provided code to visualize the inferred state sequence together with the ground truth\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:40.812495Z",
     "iopub.status.busy": "2021-06-03T14:00:40.809128Z",
     "iopub.status.idle": "2021-06-03T14:00:40.849318Z",
     "shell.execute_reply": "2021-06-03T14:00:40.848814Z"
    }
   },
   "outputs": [],
   "source": [
    "np.random.seed(101)\n",
    "nstep = 100\n",
    "switch_prob = 0.1\n",
    "log10_noise_level = -1\n",
    "\n",
    "# Build model\n",
    "model = create_HMM(switch_prob=switch_prob,\n",
    "                     noise_level=10.**log10_noise_level,\n",
    "                     startprob=[0.5, 0.5])\n",
    "observations, states = model.sample(nstep)\n",
    "# Infer state sequence\n",
    "predictive_probs, posterior_probs = simulate_forward_inference(model, nstep,\n",
    "                                                               observations)\n",
    "\n",
    "############################################################################\n",
    "# Insert your code here to:\n",
    "#      Calculate inferred states from posterior probabilities at state 0\n",
    "#      Hint: Compare the probabilities with 0.5 and note that you should\n",
    "#            return 0 if prob > 0.5\n",
    "############################################################################\n",
    "# states_inferred = ...\n",
    "\n",
    "################################################################################\n",
    "# After finishing the above exercises, please un-comment the following lines\n",
    "################################################################################\n",
    "#plot_forward_inference(model, states, observations, states_inferred)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:40.859192Z",
     "iopub.status.busy": "2021-06-03T14:00:40.858428Z",
     "iopub.status.idle": "2021-06-03T14:00:41.598802Z",
     "shell.execute_reply": "2021-06-03T14:00:41.599296Z"
    }
   },
   "source": [
    "[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D2_HiddenDynamics/solutions/W3D2_Tutorial2_Solution_c36f28a6.py)\n",
    "\n",
    "*Example output:*\n",
    "\n",
    "<img alt='Solution hint' align='left' width=560 height=416 src=https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/tutorials/W3D2_HiddenDynamics/static/W3D2_Tutorial2_Solution_c36f28a6_0.png>\n",
    "\n",
    "<img alt='Solution hint' align='left' width=560 height=416 src=https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/tutorials/W3D2_HiddenDynamics/static/W3D2_Tutorial2_Solution_c36f28a6_1.png>\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interactive Demo 4: Forward inference\n",
    "\n",
    "Try different values of switching probability (`prob_switch`) and noise level (`noise_level`) either by hand or the widget in section **Interactive Cell**. When do we start to make false inferences?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:41.632326Z",
     "iopub.status.busy": "2021-06-03T14:00:41.631700Z",
     "iopub.status.idle": "2021-06-03T14:00:42.205289Z",
     "shell.execute_reply": "2021-06-03T14:00:42.204689Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "\n",
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "np.random.seed(101)\n",
    "nstep = 100\n",
    "\n",
    "@widgets.interact\n",
    "def plot(switch_prob=(0.01, .99, .01), log10_noise_level=(-8, 1, .01)):\n",
    "  model = create_HMM(switch_prob=switch_prob,\n",
    "                      noise_level=10.**log10_noise_level,\n",
    "                      startprob = [0.5,0.5])\n",
    "  observations, states = model.sample(nstep)\n",
    "  # observations = observations.flatten()\n",
    "  # Infer state sequence\n",
    "  predictive_probs, posterior_probs = simulate_forward_inference(model, nstep, observations)\n",
    "  states_inferred = posterior_probs <= 0.5\n",
    "  plot_forward_inference(model, states, observations, states_inferred)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Summary\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Bonus"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We, the organizers, know that the next sections are much longer and more challenging than most other tutorial content. **We do not expect you to finish it—or even start it—right now**. In fact, we strongly suggest saving your time and energy for the Kalman Filtering introduced in Tutorials 3 and 4, because it will play an important role in tomorrow's material too.\n",
    "\n",
    "That said, the EM algorithm can be a very useful and powerful optimization tool. Since it is typically taught in the context of Hidden Markov Models, we have included it here for your reference. \n",
    "\n",
    "To reiterate, the remainder of this notebook is *completely* and *absolutely* optional. It is not essential to understand the rest of the NMA content. By this point in Tutorial 2, we believe that you will have seen enough about HMMs to know when/if they might be relevant for your own research. When that day comes, or you are just feeling curious, this material will be here waiting! "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Bonus Section 1: HMM for Poisson spiking neuronal network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:42.218889Z",
     "iopub.status.busy": "2021-06-03T14:00:42.218327Z",
     "iopub.status.idle": "2021-06-03T14:00:42.248911Z",
     "shell.execute_reply": "2021-06-03T14:00:42.249680Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Video 5: HMM for Poisson spiking neurons case study\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"Wb8mf5chmyI\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Given noisy neural or behavioral measurements, we as neuroscientists often want to infer the unobserved latent variables as they change over time. Thalamic relay neurons fire in two distinct modes: a tonic mode where spikes are produced one at a time, and a 'burst mode' where several action potentials are produced in rapid succession. These modes are thought to differentially encode how the neurons relay information from sensory receptors to cortex. A distinct molecular mechanism, T-type calcium channels, switches neurons between modes, but it is very challenging to measure in the brain of a living monkey. However, statistical approaches let us recover the hidden state of those calcium channels purely from their spiking activity, which can be measured in a behaving monkey.\n",
    "\n",
    "Here, we're going to tackle a simplified version of that problem.\n",
    "\n",
    "\n",
    "Let's consider the formulation mentioned in the intro lecture. \n",
    "We have a network of $C$ neurons switching between $K$ states. Neuron $c$ has firing rate $\\lambda_i^c$ in state $i$. The transition between states are represented by the $K\\times K$ transition matrix $A_{ij}$ and initial probability vector $\\psi$ with length $K$ at time $t=1$.\n",
    "\n",
    "Let $y_t^c$ be the number of spikes for cell $c$ in time bin $t$.\n",
    "\n",
    "In the following exercises (4 and 5) and tutorials, you will\n",
    "\n",
    "* Define an instance of such model with $C=5$ and $K=3$\n",
    "* (**Exercise 4**) Generate a dataset from this model\n",
    "* (**Exercise 5**) Implement the M-step for this HMM\n",
    "* Run EM to estimate all parameters $A,\\psi,\\lambda_i^c$\n",
    "* Plot the learning likelihood curve\n",
    "* Plot expected complete log likelihood versus data log likelihood\n",
    "* Compare learnt parameters versus true parameters "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "\n",
    "## **Optional** Section: Define model and generate data\n",
    "\n",
    "Let's first generate a random state sequence from the hidden Markov Chain, and generate `n_frozen_trials` different trials of spike trains for each cell assuming they all use the same underlying sequence we just generated.\n",
    "\n",
    "**Suggestions**\n",
    "\n",
    "1. Run the following two sections **Model and simulation parameters** and **Initialize true model** to define a true model and parameters that will be used in our following exercises. Please take a look at the parameters and come back to these two cells if you encounter a variable you don't know in the future.\n",
    "\n",
    "2. Complete the code to convert a given state sequence to corresponding spike rates for all cells at all times, and use provided code to visualize all spike trains.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:42.258783Z",
     "iopub.status.busy": "2021-06-03T14:00:42.252923Z",
     "iopub.status.idle": "2021-06-03T14:00:42.290111Z",
     "shell.execute_reply": "2021-06-03T14:00:42.290725Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Helper functions\n",
    "def plot_spike_train(X, Y, dt):\n",
    "    \"\"\"Plots the spike train for cells across trials and overlay the state.\n",
    "\n",
    "      Args:\n",
    "        X: (2d numpy array of binary values): The state sequence in a one-hot\n",
    "                                              representation. (T, states)\n",
    "        Y: (3d numpy array of floats):        The spike sequence.\n",
    "                                              (trials, T, C)\n",
    "        dt (float):                           Interval for a bin.\n",
    "    \"\"\"\n",
    "    n_trials, T, C = Y.shape\n",
    "    trial_T = T * dt\n",
    "    fig = plt.figure(figsize=(.7 * (12.8 + 6.4), .7 * 9.6))\n",
    "\n",
    "    # plot state sequence\n",
    "    starts = [0] + list(np.diff(X.nonzero()[1]).nonzero()[0])\n",
    "    stops = list(np.diff(X.nonzero()[1]).nonzero()[0]) + [T]\n",
    "    states = [X[i + 1].nonzero()[0][0] for i in starts]\n",
    "    for a, b, i in zip(starts, stops, states):\n",
    "        rect = patches.Rectangle((a * dt, 0), (b - a) * dt, n_trials * C,\n",
    "                                 facecolor=plt.get_cmap('tab10').colors[i],\n",
    "                                 alpha=0.15)\n",
    "        plt.gca().add_patch(rect)\n",
    "\n",
    "    # plot rasters\n",
    "    for c in range(C):\n",
    "        if c > 0:\n",
    "            plt.plot([0, trial_T], [c * n_trials, c * n_trials],\n",
    "                     color=plt.get_cmap('tab10').colors[0])\n",
    "        for r in range(n_trials):\n",
    "            tmp = Y[r, :, c].nonzero()[0]\n",
    "            if len(tmp) > 0: plt.plot(np.stack((tmp, tmp)) * dt,\n",
    "                                      (c * n_trials + r + 0.1,\n",
    "                                       c * n_trials + r + .9),\n",
    "                                      'k')\n",
    "\n",
    "    ax = plt.gca()\n",
    "    plt.yticks(np.arange(0, n_trials * C, n_trials),\n",
    "               labels=np.arange(C, dtype=int))\n",
    "    plt.xlabel('time (s)', fontsize=16)\n",
    "    plt.ylabel('Cell number', fontsize=16)\n",
    "\n",
    "def run_em(epochs, Y, psi, A, L, dt):\n",
    "  \"\"\"Run EM for the HMM spiking model.\n",
    "\n",
    "  Args:\n",
    "    epochs (int):       Number of epochs of EM to run\n",
    "    Y (numpy 3d array): Tensor of recordings, has shape (n_trials, T, C)\n",
    "    psi (numpy vector): Initial probabilities for each state\n",
    "    A (numpy matrix):   Transition matrix, A[i,j] represents the prob to switch\n",
    "                        from j to i. Has shape (K,K)\n",
    "    L (numpy matrix):   Poisson rate parameter for different cells.\n",
    "                        Has shape (C,K)\n",
    "    dt (float):         Duration of a time bin\n",
    "\n",
    "  Returns:\n",
    "    save_vals (lists of floats): Data for later plotting\n",
    "    lls (list of flots):         ll Before each EM step\n",
    "    psi (numpy vector):          Estimated initial probabilities for each state\n",
    "    A (numpy matrix):            Estimated transition matrix, A[i,j] represents\n",
    "                                 the prob to switch from j to i. Has shape (K,K)\n",
    "    L (numpy matrix):            Estimated Poisson rate parameter for different\n",
    "                                 cells. Has shape (C,K)\n",
    "  \"\"\"\n",
    "  save_vals = []\n",
    "  lls = []\n",
    "  for e in range(epochs):\n",
    "\n",
    "      # Run E-step\n",
    "      ll, gamma, xi = e_step(Y, psi, A, L, dt)\n",
    "      lls.append(ll)  # log the data log likelihood for current cycle\n",
    "\n",
    "      if e % print_every == 0: print(f'epoch: {e:3d}, ll = {ll}')  # log progress\n",
    "      # Run M-step\n",
    "      psi_new, A_new, L_new = m_step(gamma, xi, dt)\n",
    "\n",
    "      \"\"\"Booking keeping for later plotting\n",
    "      Calculate the difference of parameters for later\n",
    "      interpolation/extrapolation\n",
    "      \"\"\"\n",
    "      dp, dA, dL = psi_new - psi, A_new - A, L_new - L\n",
    "      # Calculate LLs and ECLLs for later plotting\n",
    "      if e in plot_epochs:\n",
    "          b_min = -min([np.min(psi[dp > 0] / dp[dp > 0]),\n",
    "                        np.min(A[dA > 0] / dA[dA > 0]),\n",
    "                        np.min(L[dL > 0] / dL[dL > 0])])\n",
    "          b_max = -max([np.max(psi[dp < 0] / dp[dp < 0]),\n",
    "                        np.max(A[dA < 0] / dA[dA < 0]),\n",
    "                        np.max(L[dL < 0] / dL[dL < 0])])\n",
    "          b_min = np.max([.99 * b_min, b_lims[0]])\n",
    "          b_max = np.min([.99 * b_max, b_lims[1]])\n",
    "          bs = np.linspace(b_min, b_max, num_plot_vals)\n",
    "          bs = sorted(list(set(np.hstack((bs, [0, 1])))))\n",
    "          bs = np.array(bs)\n",
    "          lls_for_plot = []\n",
    "          eclls_for_plot = []\n",
    "          for i, b in enumerate(bs):\n",
    "              ll = e_step(Y, psi + b * dp, A + b * dA, L + b * dL, dt)[0]\n",
    "              lls_for_plot.append(ll)\n",
    "              rate = (L + b * dL) * dt\n",
    "              ecll = ((gamma[:, 0] @ np.log(psi + b * dp) +\n",
    "                       (xi * np.log(A + b * dA)).sum(axis=(-1, -2, -3)) +\n",
    "                       (gamma * stats.poisson(rate).logpmf(Y[..., np.newaxis]).sum(-2)\n",
    "                       ).sum(axis=(-1, -2))).mean() / T / dt)\n",
    "              eclls_for_plot.append(ecll)\n",
    "              if b == 0:\n",
    "                  diff_ll = ll - ecll\n",
    "          lls_for_plot = np.array(lls_for_plot)\n",
    "          eclls_for_plot = np.array(eclls_for_plot) + diff_ll\n",
    "          save_vals.append((bs, lls_for_plot, eclls_for_plot))\n",
    "      # return new parameter\n",
    "      psi, A, L = psi_new, A_new, L_new\n",
    "\n",
    "  ll = e_step(Y, psi, A, L, dt)[0]\n",
    "  lls.append(ll)\n",
    "  print(f'epoch: {epochs:3d}, ll = {ll}')\n",
    "  return save_vals, lls, psi, A, L\n",
    "\n",
    "def plot_lls(lls):\n",
    "  \"\"\"Plots log likelihoods at each epoch.\n",
    "  Args:\n",
    "    lls (list of floats) log likelihoods at each epoch.\n",
    "  \"\"\"\n",
    "  epochs = len(lls)\n",
    "  fig, ax = plt.subplots()\n",
    "  ax.plot(range(epochs) , lls, linewidth=3)\n",
    "  span = max(lls) - min(lls)\n",
    "  ax.set_ylim(min(lls) - span * 0.05, max(lls) + span * 0.05)\n",
    "  plt.xlabel('iteration')\n",
    "  plt.ylabel('log likelihood')\n",
    "  plt.show(fig)\n",
    "\n",
    "def plot_lls_eclls(plot_epochs, save_vals):\n",
    "  \"\"\"Plots log likelihoods at each epoch.\n",
    "  Args:\n",
    "    plot_epochs (list of ints):  Which epochs were saved to plot.\n",
    "    save_vals (lists of floats): Different likelihoods from EM for plotting.\n",
    "  \"\"\"\n",
    "  rows = int(np.ceil(min(len(plot_epochs), len(save_vals)) / 3))\n",
    "  fig, axes = plt.subplots(rows, 3, figsize=(.7 * 6.4 * 3, .7 * 4.8 * rows))\n",
    "  axes = axes.flatten()\n",
    "\n",
    "  minll, maxll = np.inf, -np.inf\n",
    "  for i, (ax, (bs, lls_for_plot, eclls_for_plot)) in enumerate(zip(axes, save_vals)):\n",
    "      ax.set_xlim([-1.15, 2.15])\n",
    "      min_val = np.stack((lls_for_plot, eclls_for_plot)).min()\n",
    "      max_val = np.stack((lls_for_plot, eclls_for_plot)).max()\n",
    "\n",
    "      ax.plot([0, 0], [min_val, lls_for_plot[bs == 0]], '--b')\n",
    "      ax.plot([1, 1], [min_val, lls_for_plot[bs == 1]], '--b')\n",
    "      ax.set_xticks([0, 1])\n",
    "      ax.set_xticklabels([f'$\\\\theta^{plot_epochs[i]}$',\n",
    "                          f'$\\\\theta^{plot_epochs[i] + 1}$'])\n",
    "      ax.tick_params(axis='y')\n",
    "      ax.tick_params(axis='x')\n",
    "\n",
    "\n",
    "      ax.plot(bs, lls_for_plot)\n",
    "      ax.plot(bs, eclls_for_plot)\n",
    "\n",
    "      if min_val < minll: minll = min_val\n",
    "      if max_val > maxll: maxll = max_val\n",
    "\n",
    "      if i % 3 == 0: ax.set_ylabel('log likelihood')\n",
    "      if i == 4:\n",
    "          l = ax.legend(ax.lines[-2:], ['LL', 'ECLL'], framealpha=1)\n",
    "  plt.show(fig)\n",
    "\n",
    "def plot_learnt_vs_true(L_true, L, A_true, A, dt):\n",
    "  \"\"\"Plot and compare the true and learnt parameters.\n",
    "\n",
    "  Args:\n",
    "    L_true (numpy array): True L.\n",
    "    L (numpy array):      Estimated L.\n",
    "    A_true (numpy array): True A.\n",
    "    A (numpy array):      Estimated A.\n",
    "    dt (float):           Bin length.\n",
    "  \"\"\"\n",
    "  C, K = L.shape\n",
    "  fig = plt.figure(figsize=(8, 4))\n",
    "  plt.subplot(121)\n",
    "  plt.plot([0, L_true.max() * 1.05], [0, L_true.max() * 1.05], '--b')\n",
    "  for i in range(K):\n",
    "      for c in range(C):\n",
    "          plt.plot(L_true[c, i], L[c, i], color='C{}'.format(c),\n",
    "                   marker=['o', '*', 'd'][i])  # this line will fail for K > 3\n",
    "  ax = plt.gca()\n",
    "  ax.axis('equal')\n",
    "  plt.xlabel('True firing rate (Hz)')\n",
    "  plt.ylabel('Inferred firing rate (Hz)')\n",
    "  xlim, ylim = ax.get_xlim(), ax.get_ylim()\n",
    "  for c in range(C):\n",
    "      plt.plot([-10^6], [-10^6], 'o', color='C{}'.format(c))\n",
    "  for i in range(K):\n",
    "      plt.plot([-10^6], [-10^6], '.', marker=['o', '*', 'd'][i], c=\"black\")\n",
    "  l = plt.legend(ax.lines[-C - K:],\n",
    "                 [f'cell {c + 1}' for c in range(C)] + [f'state {i + 1}' for i in range(K)])\n",
    "  ax.set_xlim(xlim), ax.set_ylim(ylim)\n",
    "\n",
    "  plt.subplot(122)\n",
    "  ymax = np.max(A_true - np.diag(np.diag(A_true))) / dt * 1.05\n",
    "  plt.plot([0, ymax], [0, ymax], '--b')\n",
    "  for j in range(K):\n",
    "      for i in range(K):\n",
    "          if i == j: continue\n",
    "          plt.plot(A_true[i, j] / dt, A[i, j] / dt, 'o')\n",
    "  ax = plt.gca()\n",
    "  ax.axis('equal')\n",
    "  plt.xlabel('True transition rate (Hz)')\n",
    "  plt.ylabel('Inferred transition rate (Hz)')\n",
    "  l = plt.legend(ax.lines[1:], ['state 1 -> 2',\n",
    "                                      'state 1 -> 3',\n",
    "                                      'state 2 -> 1',\n",
    "                                      'state 2 -> 3',\n",
    "                                      'state 3 -> 1',\n",
    "                                      'state 3 -> 2'\n",
    "                                    ])\n",
    "  plt.show(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Model and simulation parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:42.296526Z",
     "iopub.status.busy": "2021-06-03T14:00:42.295196Z",
     "iopub.status.idle": "2021-06-03T14:00:42.297112Z",
     "shell.execute_reply": "2021-06-03T14:00:42.297562Z"
    }
   },
   "outputs": [],
   "source": [
    "# model and data parameters\n",
    "C = 5  # number of cells\n",
    "K = 3  # number of states\n",
    "dt = 0.002  # seconds\n",
    "trial_T = 2.0  # seconds\n",
    "n_frozen_trials = 20  # used to plot multiple trials with the same state sequence\n",
    "n_trials = 300  # number of trials (each has it's own state sequence)\n",
    "\n",
    "# for random data\n",
    "max_firing_rate = 50  # Hz\n",
    "max_transition_rate = 3  # Hz\n",
    "\n",
    "# needed to plot LL and ECLL for every M-step\n",
    "# **This substantially slows things down!!**\n",
    "num_plot_vals = 10  # resolution of the plot (this is the expensive part)\n",
    "b_lims = (-1, 2)  # lower limit on graph (b = 0 is start-of-M-step LL; b = 1 is end-of-M-step LL)\n",
    "plot_epochs = list(range(9))  # list of epochs to plot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Initialize true model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:42.305109Z",
     "iopub.status.busy": "2021-06-03T14:00:42.303741Z",
     "iopub.status.idle": "2021-06-03T14:00:42.305730Z",
     "shell.execute_reply": "2021-06-03T14:00:42.306205Z"
    }
   },
   "outputs": [],
   "source": [
    "np.random.seed(101)\n",
    "T = round(trial_T / dt)\n",
    "ts = np.arange(T)\n",
    "\n",
    "# initial state distribution\n",
    "psi = np.arange(1, K + 1)\n",
    "psi = psi / psi.sum()\n",
    "\n",
    "# off-diagonal transition rates sampled uniformly\n",
    "A = np.random.rand(K, K) * max_transition_rate * dt\n",
    "A = (1. - np.eye(K)) * A\n",
    "A = A + np.diag(1 - A.sum(1))\n",
    "\n",
    "# hand-crafted firing rates make good plots\n",
    "L = np.array([\n",
    "    [.02, .8, .37],\n",
    "    [1., .7, .1],\n",
    "    [.92, .07, .5],\n",
    "    [.25, .42, .75],\n",
    "    [.15, .2, .85]\n",
    "]) * max_firing_rate     # (C,K)\n",
    "\n",
    "# Save true parameters for comparison later\n",
    "psi_true = psi\n",
    "A_true = A\n",
    "L_true = L"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Generate data with frozen sequence and plot\n",
    "Given a state sequence `[0,1,1,3,2,...]`, we'll first convert each state in to sequence in to the so-called \"one-hot\" coding. For example, with 5 total states, the one-hot coding of state `0` is `[1,0,0,0,0]` and the coding for state `3` is `[0,0,0,1,0]`. Suppose we now have a sequence of length `T`, the one-hot coding of this sequence `Xf` will have shape `(T,K)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:42.316960Z",
     "iopub.status.busy": "2021-06-03T14:00:42.315973Z",
     "iopub.status.idle": "2021-06-03T14:00:50.760481Z",
     "shell.execute_reply": "2021-06-03T14:00:50.760942Z"
    }
   },
   "outputs": [],
   "source": [
    "np.random.seed(101)\n",
    "# sample n_frozen_trials state sequences\n",
    "Xf = np.zeros(T, dtype=int)\n",
    "Xf[0] = (psi.cumsum() > np.random.rand()).argmax()\n",
    "for t in range(1, T):\n",
    "    Xf[t] = (A[Xf[t - 1],:].cumsum() > np.random.rand()).argmax()\n",
    "\n",
    "# switch to one-hot encoding of the state\n",
    "Xf = np.eye(K, dtype=int)[Xf]  # (T,K)\n",
    "\n",
    "# get the Y values\n",
    "Rates = np.squeeze(L @ Xf[..., None]) * dt  # (T,C)\n",
    "\n",
    "Rates = np.tile(Rates, [n_frozen_trials, 1, 1]) # (n_trials, T, C)\n",
    "Yf = stats.poisson(Rates).rvs()\n",
    "\n",
    "with plt.xkcd():\n",
    "  plot_spike_train(Xf, Yf, dt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Generate data for EM learning\n",
    "\n",
    "The previous dataset is generated with the same state sequence for visualization. Now let's generate `n_trials` trials of observations, each one with its own randomly generated sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:50.767754Z",
     "iopub.status.busy": "2021-06-03T14:00:50.767157Z",
     "iopub.status.idle": "2021-06-03T14:00:50.884250Z",
     "shell.execute_reply": "2021-06-03T14:00:50.884722Z"
    }
   },
   "outputs": [],
   "source": [
    "np.random.seed(101)\n",
    "# sample n_trials state sequences\n",
    "X = np.zeros((n_trials, T), dtype=int)\n",
    "X[:, 0] = (psi_true.cumsum(0)[:, None] > np.random.rand(n_trials)).argmax(0)\n",
    "for t in range(1, T):\n",
    "  X[:, t] = (A_true[X[:, t - 1], :].T.cumsum(0) > np.random.rand(n_trials)).argmax(0)\n",
    "\n",
    "# switch to one-hot encoding of the state\n",
    "one_hot = np.eye(K)[np.array(X).reshape(-1)]\n",
    "X = one_hot.reshape(list(X.shape) + [K])\n",
    "\n",
    "# get the Y values\n",
    "Y = stats.poisson(np.squeeze(L_true @ X[..., None]) * dt).rvs()  # (n_trials, T, C)\n",
    "print(\"Y has shape: (n_trial={},T={},C={})\".format(*Y.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## **Optional** Section: EM algorithm for HMM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:50.890767Z",
     "iopub.status.busy": "2021-06-03T14:00:50.890114Z",
     "iopub.status.idle": "2021-06-03T14:00:50.925272Z",
     "shell.execute_reply": "2021-06-03T14:00:50.924716Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Video 6: EM Tutorial\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"umU4wUWlKvg\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finding the optimal values of parameters that maximizes the data likelihood is practically infeasible since we need to integrating out all latent variables $x_{1:T}$. The time needed is exponential to $T$. Thus as an alternative approach, we use the Expectation-Maximization algorithm, which iteratively performing a E-step followed by a M-step and is guaranteed to not decrease(usually increase) the data likelihood after each EM cycle.\n",
    "\n",
    "\n",
    "In this section we will briefly review the EM algorithm  for HMM and list\n",
    "\n",
    "* Recursive equations for forward and backward probabilities $a_i(t)$ and $b_i(t)$\n",
    "* Expressions for singleton and pairwise marginal distributions after seeing data: $\\gamma_{i}(t):=p_{\\theta}\\left(x_{t}=i | Y_{1: T}\\right)$ and $\\xi_{i j}(t) = p_{\\theta}(x_t=i,x_{t+1}=j|Y_{1:T})$\n",
    "* Closed-form solutions for updated values of $A,\\psi,\\lambda$ which increases data likelihood\n",
    "\n",
    "\n",
    "---\n",
    "#### E-step: Forward-backward algorithm \n",
    "In the forward pass, we calculate the **forward probabilities**, or the joint probability of $x_t$ and current and past data $Y_{1:t}$: $a_i(t):=p(x_t=i,Y_{1:t})$ recursively by\n",
    "\n",
    "$$a_i(t) = p_(y_t|x_i=t)\\sum_j A_{ji} a_j(t-1)$$\n",
    "\n",
    "In contrast to the intro, now $A_{ji}$ means **the transition probability from state $j$ to state $i$.**\n",
    "\n",
    "The backward pass calculate the **backward probabilities** $b_i(t):=p_{\\theta}(Y_{t+1:T}|x_t=i)$, which is the likelihood of observing all future data points given current state $x_t$. The recursion of $b_i(t)$ is given by \n",
    "\n",
    "$$ b_i(t) = \\sum_j p_{\\theta}(y_{t+1}|x_{t+1}=j)b_j(t+1)A_{ij} $$\n",
    "\n",
    "Combining all past and future information, the **singleton and pairwise marginal distributions** are given by \n",
    "\n",
    "$$ \\gamma_{i}(t):=p_{\\theta}\\left(x_{t}=i | Y_{1: T}\\right)=\\frac{a_{i}(t) b_{i}(t)}{p_{\\theta}\\left(Y_{1: T}\\right)} $$ \n",
    "\n",
    "$$  \\xi_{i j}(t) = p_{\\theta}(x_t=i,x_{t+1}=j|Y_{1:T}) =\\frac{b_{j}(t+1)p_{\\theta}\\left(y_{t+1} | x_{t+1}=j\\right) A_{i j} a_{i}(t)}{p_{\\theta}\\left(Y_{1: T}\\right)} $$\n",
    "\n",
    "where $p_{\\theta}(Y_{1:T})=\\sum_i a_i(T)$.\n",
    "\n",
    "---\n",
    "#### M-step\n",
    "\n",
    "The M-step for HMM has a closed-form solution. First the new transition matrix is given by \n",
    "$$ \n",
    " A_{ij} =\\frac{\\sum_{t=1}^{T-1} \\xi_{i j}(t)}{\\sum_{t=1}^{T-1} \\gamma_{i}(t)}\n",
    "$$\n",
    "\n",
    "which is the expected empirical transition probabilities. \n",
    "New initial probabilities and parameters of the emission models are also given by their empirical values given single and pairwise marginal distributions:\n",
    "\n",
    "$$ \\psi_i = \\frac{1}{N}\\sum_{trials}\\gamma_i(1) $$ \n",
    "\n",
    "$$ \\lambda_{i}^{c}=\\frac{\\sum_{t} \\gamma_{i}(t) y_{t}^{c}}{\\sum_{t} \\gamma_{i}(t) d t}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "\n",
    "### E-step: forward and backward algorithm\n",
    "\n",
    "**(Optional)**\n",
    "\n",
    "In this section you will read through the code for the forward-backward algorithm and understand how to implement the computation efficiently in `numpy` by calculating the recursion for all trials at once.\n",
    "\n",
    "---\n",
    "\n",
    "Let's re-write the forward and backward recursions in a more compact form:\n",
    "\n",
    "$$ a_i^t = \\sum_j A_{ji}o_j^t a_j^{t-1}  $$\n",
    "\n",
    "\n",
    "$$b^t_i = \\sum_j A_{ij} o_j^{t+1}b_j^{t+1} $$ where $o_j^{t}=p(y_{t}|x_{t}=j)$.\n",
    "\n",
    "\n",
    "Let's take the backward recursion for example. In practice we will handle all trials together since they are independent of each other. After adding a trial index $l$ to the recursion equations, the backward recursion becomes:\n",
    "\n",
    "$$b^t_{li} = \\sum_j A_{ij} o_{lj}^{t+1}b_{lj}^{t+1} $$ \n",
    "\n",
    "What we have in hand are:\n",
    "* `A`: matrix of size `(K,K)`\n",
    "* `o^{t+1}`: array of size `(N,K)` is the log data likelihood for all trials at a given time\n",
    "* `b^{t+1}`: array of size `(N,K)` is the backward probability for all trials at a given time\n",
    "\n",
    "where `N` stands for the number of trials.\n",
    "\n",
    "The index size and meaning doesn't match for these three arrays: the index is $i$ for $A$ in the first dimension and is $l$ for $o$ and $b$, so we can't just multiply them together. However, we can do this by viewing vectors $o^{t+1}_{l\\cdot}$ and $b^{t+1}_{l\\cdot}$ as a matrix with 1 row and re-write the backward equation as:\n",
    "\n",
    "$$b^t_{li} = \\sum_j A_{ij} o_{l1j}^{t+1}b_{l1j}^{t+1} $$\n",
    "\n",
    "Now we can just multiply these three arrays element-wise and sum over the last dimension.\n",
    "\n",
    "In `numpy`, we can achieve this by indexing the array with `None` at the location we want to insert a dimension. Take `b` with size `(N,T,K)` for example,`b[:,t,:]` will have shape `(N,K)`, `b[:,t,None,:]` will have shape `(N,1,K)` and `b[:,t,:,None]` will have shape `(N,K,1)`.\n",
    "\n",
    "So the backward recursion computation can be implemented as \n",
    "\n",
    "```python\n",
    "b[:,t,:] = (A * o[:,t+1,None,:] * b[:,t+1,None,:]).sum(-1)\n",
    "```\n",
    "\n",
    "---\n",
    "\n",
    "In addition to the trick introduced above, in this exercise we will work in the **log scale** for numerical stability.\n",
    "\n",
    "\n",
    "**Suggestions** \n",
    "\n",
    "1. Take a look at the code for the forward recursion and backward recursion.\n",
    "\n",
    "    \n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:50.938219Z",
     "iopub.status.busy": "2021-06-03T14:00:50.936486Z",
     "iopub.status.idle": "2021-06-03T14:00:50.939019Z",
     "shell.execute_reply": "2021-06-03T14:00:50.939781Z"
    }
   },
   "outputs": [],
   "source": [
    "def e_step(Y, psi, A, L, dt):\n",
    "  \"\"\"Calculate the E-step for the HMM spiking model.\n",
    "\n",
    "  Args:\n",
    "    Y (numpy 3d array): tensor of recordings, has shape (n_trials, T, C)\n",
    "    psi (numpy vector): initial probabilities for each state\n",
    "    A (numpy matrix):   transition matrix, A[i,j] represents the prob to\n",
    "                        switch from i to j. Has shape (K,K)\n",
    "    L (numpy matrix):   Poisson rate parameter for different cells.\n",
    "                        Has shape (C,K)\n",
    "    dt (float):         Bin length\n",
    "\n",
    "  Returns:\n",
    "    ll (float):             data log likelihood\n",
    "    gamma (numpy 3d array): singleton marginal distribution.\n",
    "                            Has shape (n_trials, T, K)\n",
    "    xi (numpy 4d array):    pairwise marginal distribution for adjacent\n",
    "                            nodes . Has shape (n_trials, T-1, K, K)\n",
    "  \"\"\"\n",
    "  n_trials = Y.shape[0]\n",
    "  T = Y.shape[1]\n",
    "  K = psi.size\n",
    "  log_a = np.zeros((n_trials, T, K))\n",
    "  log_b = np.zeros((n_trials, T, K))\n",
    "\n",
    "  log_A = np.log(A)\n",
    "  log_obs = stats.poisson(L * dt).logpmf(Y[..., None]).sum(-2)  # n_trials, T, K\n",
    "\n",
    "  # forward pass\n",
    "  log_a[:, 0] = log_obs[:, 0] + np.log(psi)\n",
    "  for t in range(1, T):\n",
    "    tmp = log_A + log_a[:, t - 1, : ,None]  # (n_trials, K,K)\n",
    "    maxtmp = tmp.max(-2)  # (n_trials,K)\n",
    "    log_a[:, t] = (log_obs[:, t] + maxtmp +\n",
    "                    np.log(np.exp(tmp - maxtmp[:, None]).sum(-2)))\n",
    "\n",
    "  # backward pass\n",
    "  for t in range(T - 2, -1, -1):\n",
    "    tmp = log_A + log_b[:, t + 1, None] + log_obs[:, t + 1, None]\n",
    "    maxtmp = tmp.max(-1)\n",
    "    log_b[:, t] = maxtmp + np.log(np.exp(tmp - maxtmp[..., None]).sum(-1))\n",
    "\n",
    "  # data log likelihood\n",
    "  maxtmp = log_a[:, -1].max(-1)\n",
    "  ll = np.log(np.exp(log_a[:, -1] - maxtmp[:, None]).sum(-1)) + maxtmp\n",
    "\n",
    "  # singleton and pairwise marginal distributions\n",
    "  gamma = np.exp(log_a + log_b - ll[:, None, None])\n",
    "  xi = np.exp(log_a[:, :-1, :, None] + (log_obs + log_b)[:, 1:, None] +\n",
    "              log_A - ll[:, None, None, None])\n",
    "\n",
    "  return ll.mean() / T / dt, gamma, xi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:50.945530Z",
     "iopub.status.busy": "2021-06-03T14:00:50.944688Z",
     "iopub.status.idle": "2021-06-03T14:00:50.976579Z",
     "shell.execute_reply": "2021-06-03T14:00:50.977152Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title EXERCISE 7: Implement the M-step Video\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"H4GGTg_9BaE\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "### EXERCISE 5: Implement the M-step\n",
    "\n",
    "In this exercise you will complete the M-step for this HMM using closed form solutions mentioned before.\n",
    "\n",
    "**Suggestions**\n",
    "\n",
    "1. Calculate new initial probabilities as empirical counts of singleton marginals\n",
    "\n",
    "$$ \\psi_i = \\frac{1}{N}\\sum_{trials}\\gamma_i(1) $$\n",
    "\n",
    "2. Remember the extra trial dimension and average over all trials\n",
    "\n",
    "\n",
    "**For reference:**\n",
    "\n",
    "New transition matrix is calculated as empirical counts of transition events from marginals\n",
    "\n",
    "$$ A_{ij} =\\frac{\\sum_{t=1}^{T-1} \\xi_{i j}(t)}{\\sum_{t=1}^{T-1} \\gamma_{i}(t)}$$ \n",
    "\n",
    "\n",
    "New spiking rates for each cell and each state are given by\n",
    "\n",
    "$$  \\lambda_{i}^{c}=\\frac{\\sum_{t} \\gamma_{i}(t) y_{t}^{c}}{\\sum_{t} \\gamma_{i}(t) d t} $$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:50.985333Z",
     "iopub.status.busy": "2021-06-03T14:00:50.983774Z",
     "iopub.status.idle": "2021-06-03T14:00:50.986091Z",
     "shell.execute_reply": "2021-06-03T14:00:50.986610Z"
    }
   },
   "outputs": [],
   "source": [
    "def m_step(gamma, xi, dt):\n",
    "  \"\"\"Calculate the M-step updates for the HMM spiking model.\n",
    "\n",
    "  Args:\n",
    "    gamma ():       Number of epochs of EM to run\n",
    "    xi (numpy 3d array): Tensor of recordings, has shape (n_trials, T, C)\n",
    "    dt (float):         Duration of a time bin\n",
    "\n",
    "  Returns:\n",
    "    psi_new (numpy vector): Updated initial probabilities for each state\n",
    "    A_new (numpy matrix):   Updated transition matrix, A[i,j] represents the\n",
    "                            prob. to switch from j to i. Has shape (K,K)\n",
    "    L_new (numpy matrix):   Updated Poisson rate parameter for different\n",
    "                            cells. Has shape (C,K)\n",
    "  \"\"\"\n",
    "  raise NotImplementedError(\"`m_step` need to be implemented\")\n",
    "  ############################################################################\n",
    "  # Insert your code here to:\n",
    "  #    Calculate the new prior probabilities in each state at time 0\n",
    "  #    Hint: Take the first time step and average over all trials\n",
    "  ###########################################################################\n",
    "  psi_new = ...\n",
    "  # Make sure the probabilities are normalized\n",
    "  psi_new /= psi_new.sum()\n",
    "\n",
    "  # Calculate new transition matrix\n",
    "  A_new = xi.sum(axis=(0, 1)) / gamma[:, :-1].sum(axis=(0, 1))[:, np.newaxis]\n",
    "  # Calculate new firing rates\n",
    "  L_new = (np.swapaxes(Y, -1, -2) @ gamma).sum(axis=0) / gamma.sum(axis=(0, 1)) / dt\n",
    "  return psi_new, A_new, L_new"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:50.993574Z",
     "iopub.status.busy": "2021-06-03T14:00:50.992274Z",
     "iopub.status.idle": "2021-06-03T14:00:50.994325Z",
     "shell.execute_reply": "2021-06-03T14:00:50.994925Z"
    }
   },
   "source": [
    "[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D2_HiddenDynamics/solutions/W3D2_Tutorial2_Solution_ab737584.py)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:51.001085Z",
     "iopub.status.busy": "2021-06-03T14:00:51.000432Z",
     "iopub.status.idle": "2021-06-03T14:00:51.033800Z",
     "shell.execute_reply": "2021-06-03T14:00:51.034436Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Video 8: Running and plotting EM\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"6UTsXxE3hG0\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "### Run EM\n",
    "\n",
    "####Initialization for parameters\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:51.041066Z",
     "iopub.status.busy": "2021-06-03T14:00:51.040407Z",
     "iopub.status.idle": "2021-06-03T14:00:51.044128Z",
     "shell.execute_reply": "2021-06-03T14:00:51.044635Z"
    }
   },
   "outputs": [],
   "source": [
    "np.random.seed(101)\n",
    "# number of EM steps\n",
    "epochs = 9\n",
    "print_every = 1\n",
    "\n",
    "# initial state distribution\n",
    "psi = np.arange(1, K + 1)\n",
    "psi = psi / psi.sum()\n",
    "\n",
    "# off-diagonal transition rates sampled uniformly\n",
    "A = np.ones((K, K)) * max_transition_rate * dt / 2\n",
    "A = (1 - np.eye(K)) * A\n",
    "A = A + np.diag(1 - A.sum(1))\n",
    "\n",
    "# firing rates sampled uniformly\n",
    "L = np.random.rand(C, K) * max_firing_rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:00:51.049271Z",
     "iopub.status.busy": "2021-06-03T14:00:51.048537Z",
     "iopub.status.idle": "2021-06-03T14:03:07.290384Z",
     "shell.execute_reply": "2021-06-03T14:03:07.290833Z"
    }
   },
   "outputs": [],
   "source": [
    "# LL for true vs. initial parameters\n",
    "print(f'LL for true 𝜃:    {e_step(Y, psi_true, A_true, L_true, dt)[0]}')\n",
    "print(f'LL for initial 𝜃: {e_step(Y, psi, A, L, dt)[0]}\\n')\n",
    "\n",
    "# Run EM\n",
    "save_vals, lls, psi, A, L = run_em(epochs, Y, psi, A, L, dt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:03:07.297821Z",
     "iopub.status.busy": "2021-06-03T14:03:07.296511Z",
     "iopub.status.idle": "2021-06-03T14:03:07.298511Z",
     "shell.execute_reply": "2021-06-03T14:03:07.298981Z"
    }
   },
   "outputs": [],
   "source": [
    "# EM doesn't guarantee the order of learnt latent states are the same as that of true model\n",
    "# so we need to sort learnt parameters\n",
    "\n",
    "# Compare all true and estimated latents across cells\n",
    "cost_mat = np.sum((L_true[..., np.newaxis] - L[:, np.newaxis])**2, axis=0)\n",
    "true_ind, est_ind = linear_sum_assignment(cost_mat)\n",
    "\n",
    "psi = psi[est_ind]\n",
    "A = A[est_ind]\n",
    "A = A[:, est_ind]\n",
    "L = L[:, est_ind]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## **Optional** Section: Plotting the training process and learnt model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting progress during EM!\n",
    "\n",
    "Now you can\n",
    "\n",
    "* Plot the likelihood during training\n",
    "* Plot the M-step log likelihood versus expected complete log likelihood(ECLL) to get an intuition of how EM works and the convexity of ECLL\n",
    "* Plot learnt parameters versus true parameters "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:03:07.327176Z",
     "iopub.status.busy": "2021-06-03T14:03:07.326079Z",
     "iopub.status.idle": "2021-06-03T14:03:07.554462Z",
     "shell.execute_reply": "2021-06-03T14:03:07.555000Z"
    }
   },
   "outputs": [],
   "source": [
    "# Plot the log likelihood after each epoch of EM\n",
    "with plt.xkcd():\n",
    "  plot_lls(lls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:03:07.649097Z",
     "iopub.status.busy": "2021-06-03T14:03:07.600144Z",
     "iopub.status.idle": "2021-06-03T14:03:09.738381Z",
     "shell.execute_reply": "2021-06-03T14:03:09.738889Z"
    }
   },
   "outputs": [],
   "source": [
    "# For each saved epoch, plot the log likelihood and expected complete log likelihood\n",
    "# for the initial and final parameter values\n",
    "with plt.xkcd():\n",
    "  plot_lls_eclls(plot_epochs, save_vals)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot learnt parameters vs. true parameters\n",
    "\n",
    "Now we will plot the (sorted) learnt parameters with true parameters to see if we successfully recovered all the parameters\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:03:09.827111Z",
     "iopub.status.busy": "2021-06-03T14:03:09.795433Z",
     "iopub.status.idle": "2021-06-03T14:03:10.479188Z",
     "shell.execute_reply": "2021-06-03T14:03:10.479736Z"
    }
   },
   "outputs": [],
   "source": [
    "# Compare true and learnt parameters\n",
    "with plt.xkcd():\n",
    "  plot_learnt_vs_true(L_true, L, A_true, A, dt)"
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "W3D2_Tutorial2",
   "provenance": [],
   "toc_visible": true
  },
  "kernel": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": true,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
